package comlib.table

import chisel3._
import chisel3.util._

class TableCell[T<:Data](gen:T, rnum:Int=1, wnum:Int=1) extends Module{

  class Write[T<:Data](gen:T) extends Bundle{
    val en_row = Input(Bool())
    val en_col = Input(Bool())
    val data = Input(chiselTypeOf(gen))
    override def cloneType = new Write(gen).asInstanceOf[this.type]
  }

  class Read[T<:Data](gen:T) extends Bundle{
    val en_row = Input(Bool())
    val en_col = Input(Bool())
    val data = Output(chiselTypeOf(gen))
    override def cloneType = new Read(gen).asInstanceOf[this.type]
  }

  class TableCellIO extends Bundle{
    val write = Vec(wnum, new Write(gen))
    val read = Vec(rnum, new Read(gen))
  }

  val io = IO(new TableCellIO)

  val data = RegInit(gen)

  for( r <- 0 until rnum ){
    io.read(r).data := Mux( io.read(r).en_row & io.read(r).en_col, data, gen )  
  }

  val wdata = io.write.map{ case w => Mux(w.en_row & w.en_col, w.data, gen) }.map(_.asUInt).reduce(_|_).asTypeOf(gen)

  data := wdata
}

class TableOH[T<:Data](gen:T, ncol:Int, nrow:Int=1, rnum:Int=1, wnum:Int=1) extends MultiIOModule{

  class Write[T<:Data](gen:T) extends Bundle{
    val row = Input(UInt(nrow.W))
    val col = Input(UInt(ncol.W))
    val data = Input(chiselTypeOf(gen))
    override def cloneType = new Write(gen).asInstanceOf[this.type]
  }

  class Read[T<:Data](gen:T) extends Bundle{
    val row = Input(UInt(nrow.W))
    val col = Input(UInt(ncol.W))
    val data = Output(chiselTypeOf(gen))  
    override def cloneType = new Read(gen).asInstanceOf[this.type]
  }

  val readIO = Vec(rnum, new Read(gen))
  val writeIO = Vec(wnum, new Write(gen))
  val cell = Seq.fill(nrow)(Seq.fill(ncol)(Module(new TableCell[T](gen, rnum, wnum))))

  for( row <- 0 until nrow ){
    for( col <- 0 until ncol ){

      for( rd <- 0 until rnum ){
        cell(row)(col).io.read(rd).en_row := readIO(rd).row(row)
        cell(row)(col).io.read(rd).en_col := readIO(rd).col(row)
      }  

      for( wr <- 0 until wnum ){
        cell(row)(col).io.write(wr).en_row := writeIO(wr).row(row)  
        cell(row)(col).io.write(wr).en_col := writeIO(wr).col(row)
        cell(row)(col).io.write(wr).data := writeIO(wr).data  
      }
    }  
  }

  for( rd <- 0 until rnum ){
    val data_row = cell.map { cell_row => cell_row.map(_.io.read(rd).data.asUInt).reduce(_|_) }
    readIO(rd).data := data_row.reduce(_|_).asTypeOf(gen)
  }

}